#!/usr/bin/python
import hmmtop
import ancescon
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from Bio.Alphabet import IUPAC

class split:

	def __init__(self):
		self.flank = False
		self.loops = [10,10]
		self.tms = False
		self.fasta_file = False
		self.outfile = False
		self.loops = {}
		
	def __call__(self):
		hmt = hmmtop.tools()
		hmt.add_library('TMT',self.fasta_file)
		hmt.scan_libraries()
		self.tms=hmt.results['TMT']
		self.build_flanks()
		self.write_file()
	
	def build_flanks(self):
		fastas = SeqIO.to_dict(SeqIO.parse(self.fasta_file,'fasta'))
		for symbol, tmss in self.tms.items():
			for tms, residues in tmss.items():
				prev = tmss[tms-1] if tmss.has_key(tms-1) else False
				next = tmss[tms+1] if tmss.has_key(tms+1) else False
				this = self.expand(prev,residues,next,len(fastas[symbol].seq))
				loop = fastas[symbol].seq[this[0]-1:this[1]]
				self.loops.setdefault(symbol,[]).append(loop)				
				
	def expand(self,tmp,tmt,tmn,length):
		if tmp is False:
			# This is the first tms
			if self.flank is False:
				# Use the whole len, no flank specified
				tmt[0] = 1
			else:
				if tmt[0] - self.flank[0] < 1:
					tmt[0] = 1
				else:
					tmt[0] -= self.flank[0]
		else:
			# Not the first TMS, getting new start
			if self.flank is False:
				# No left flank specified, find midpoint
				left = (tmp[1]+tmt[0])
				if left%2==0:
					# Its even, use this value
					tmt[0]=left/2
				else:
					# Its odd, round up
					tmt[0]=int(round(left/2))
			else:
				# Flanking zones are specified
				tmt[0] -= self.flank[0]
				
		# Done With start position, now calculate end position
		if tmn is False:
			# This is the last tms position
			if self.flank is False:
				# Use the remaining len of sequenec
				tmt[1] = length
			else:
				if tmt[1] + self.flank[1] > length:
					tmt[1] = length
				else:
					tmt[1] += self.flank[1]
		else:
			# This is not the last TMS, getting new END position
			if self.flank is False:
				# No right flank, find midpoint, round DOWN
				right = (tmt[1]+tmn[0])
				if right%2==0:
					tmt[1] = (right/2)-1
				else:
					tmt[1] = int(right/2)
			else:
				# Flanking zones are specified
				tmt[1] += self.flank[1]
		return tmt
	
	def write_file(self):
		handle = open(self.outfile,"wb+")
		for symbol,tmss in self.loops.items():
			line = [symbol]
			seq = [str(i) for i in tmss]
			line.extend(seq)
			handle.write("\t".join(line)+"\n")
			
		
if __name__=='__main__':
	import sys
	sp = split()
	sp.fasta_file = sys.argv[1]
	sp.outfile = sys.argv[2]
	sp()